// ============================================================================
// File: pkg/snapshotter/snapshotter.go
// ============================================================================
package snapshotter

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/snapshots"
	"github.com/sirupsen/logrus"

	"gitee.com/wangtsingx/onyx-snapshotter/pkg/cache"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/config"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/content"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/lazy"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/metadata"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/metrics"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/prefetch"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/puller"
)

type Snapshotter struct {
	root      string
	ms        *metadata.Store
	cs        *content.Store
	cache     *cache.LRUCache
	lazy      *lazy.Loader
	puller    *puller.Puller
	prefetch  *prefetch.Engine
	cfg       *config.Config
	mu        sync.RWMutex
}

func NewSnapshotter(cfg *config.Config) (*Snapshotter, error) {
	// Create root directories
	dirs := []string{
		cfg.Root,
		filepath.Join(cfg.Root, "content"),
		filepath.Join(cfg.Root, "snapshots"),
		filepath.Join(cfg.Root, "cache"),
	}
	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("failed to create dir %s: %w", dir, err)
		}
	}

	// Initialize metadata store
	ms, err := metadata.NewStore(filepath.Join(cfg.Root, "metadata.db"))
	if err != nil {
		return nil, fmt.Errorf("failed to create metadata store: %w", err)
	}

	// Initialize content store with block-level deduplication
	cs, err := content.NewStore(filepath.Join(cfg.Root, "content"), cfg.Storage.BlockSize)
	if err != nil {
		return nil, fmt.Errorf("failed to create content store: %w", err)
	}

	// Initialize LRU cache
	cacheSize := parseSize(cfg.Storage.CacheSize)
	lruCache := cache.NewLRUCache(cacheSize)

	// Initialize puller with concurrent download support
	p := puller.NewPuller(&puller.Config{
		RegistryURL:   cfg.Registry.URL,
		Username:      cfg.Registry.Auth.Username,
		Password:      cfg.Registry.Auth.Password,
		Timeout:       cfg.Registry.Timeout,
		MaxRetries:    cfg.Registry.MaxRetries,
		MaxConcurrent: cfg.Storage.MaxConcurrent,
	}, cs, lruCache)

	// Initialize prefetch engine
	var pf *prefetch.Engine
	if cfg.Lazy.Prefetch.Enabled {
		pf = prefetch.NewEngine(&prefetch.Config{
			Paths:          cfg.Lazy.Prefetch.Paths,
			MaxSize:        parseSize(cfg.Lazy.Prefetch.MaxSize),
			ConcurrentJobs: cfg.Lazy.Prefetch.ConcurrentJobs,
		}, cs, p)
	}

	// Initialize lazy loader
	var loader *lazy.Loader
	if cfg.Lazy.Enabled {
		loader, err = lazy.NewLoader(cfg.Root, cs, p, lruCache)
		if err != nil {
			return nil, fmt.Errorf("failed to create lazy loader: %w", err)
		}
	}

	s := &Snapshotter{
		root:     cfg.Root,
		ms:       ms,
		cs:       cs,
		cache:    lruCache,
		lazy:     loader,
		puller:   p,
		prefetch: pf,
		cfg:      cfg,
	}

	// Start lazy loader
	if loader != nil {
		go func() {
			if err := loader.Start(context.Background()); err != nil {
				logrus.WithError(err).Error("lazy loader stopped")
			}
		}()
	}

	logrus.Info("snapshotter initialized successfully")
	return s, nil
}

func (s *Snapshotter) Stat(ctx context.Context, key string) (snapshots.Info, error) {
	metrics.SnapshotOperations.WithLabelValues("stat").Inc()
	
	s.mu.RLock()
	defer s.mu.RUnlock()

	info, err := s.ms.Get(ctx, key)
	if err != nil {
		metrics.SnapshotErrors.WithLabelValues("stat").Inc()
		return snapshots.Info{}, err
	}

	return snapshots.Info{
		Name:   info.Name,
		Parent: info.Parent,
		Kind:   info.Kind,
		Labels: info.Labels,
	}, nil
}

func (s *Snapshotter) Update(ctx context.Context, info snapshots.Info, fieldpaths ...string) (snapshots.Info, error) {
	metrics.SnapshotOperations.WithLabelValues("update").Inc()
	
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.ms.Update(ctx, info, fieldpaths...)
}

func (s *Snapshotter) Usage(ctx context.Context, key string) (snapshots.Usage, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Check if snapshot exists
	if _, err := s.ms.Get(ctx, key); err != nil {
		return snapshots.Usage{}, err
	}

	// Since snapshots.Info doesn't have Size and Inodes fields in containerd API,
	// we return zero values for now.
	return snapshots.Usage{
		Size:   0,
		Inodes: 0,
	}, nil
}

func (s *Snapshotter) Mounts(ctx context.Context, key string) ([]mount.Mount, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	info, err := s.ms.Get(ctx, key)
	if err != nil {
		return nil, err
	}

	return s.mounts(info, key)
}

func (s *Snapshotter) Prepare(ctx context.Context, key, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
	metrics.SnapshotOperations.WithLabelValues("prepare").Inc()
	timer := metrics.NewTimer()
	defer func() {
		metrics.SnapshotDuration.WithLabelValues("prepare").Observe(timer.ObserveDuration().Seconds())
	}()

	s.mu.Lock()
	defer s.mu.Unlock()

	return s.createSnapshot(ctx, snapshots.KindActive, key, parent, opts...)
}

func (s *Snapshotter) View(ctx context.Context, key, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
	metrics.SnapshotOperations.WithLabelValues("view").Inc()
	
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.createSnapshot(ctx, snapshots.KindView, key, parent, opts...)
}

func (s *Snapshotter) Commit(ctx context.Context, name, key string, opts ...snapshots.Opt) error {
	metrics.SnapshotOperations.WithLabelValues("commit").Inc()
	
	s.mu.Lock()
	defer s.mu.Unlock()

	info, err := s.ms.Get(ctx, key)
	if err != nil {
		metrics.SnapshotErrors.WithLabelValues("commit").Inc()
		return err
	}

	if info.Kind != snapshots.KindActive {
		return fmt.Errorf("snapshot %s is not active", key)
	}

	return s.ms.Commit(ctx, key, name, opts...)
}

func (s *Snapshotter) Remove(ctx context.Context, key string) error {
	metrics.SnapshotOperations.WithLabelValues("remove").Inc()
	
	s.mu.Lock()
	defer s.mu.Unlock()

	_, err := s.ms.Get(ctx, key)
	if err != nil {
		return err
	}

	// Remove snapshot directory
	snapDir := filepath.Join(s.root, "snapshots", key) // Use key instead of info.ID
	if err := os.RemoveAll(snapDir); err != nil {
		logrus.WithError(err).Warnf("failed to remove snapshot dir %s", snapDir)
	}

	// Update metrics
	metrics.ActiveSnapshots.Dec()

	return s.ms.Remove(ctx, key)
}

func (s *Snapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, fs ...string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.ms.Walk(ctx, fn, fs...)
}

func (s *Snapshotter) Close() error {
	logrus.Info("closing snapshotter")
	
	if s.lazy != nil {
		s.lazy.Stop()
	}
	if s.prefetch != nil {
		s.prefetch.Stop()
	}
	
	return s.ms.Close()
}

func (s *Snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, key, parent string, opts ...snapshots.Opt) ([]mount.Mount, error) {
	// Create snapshot metadata
	info, err := s.ms.Create(ctx, kind, key, parent, opts...)
	if err != nil {
		metrics.SnapshotErrors.WithLabelValues("create").Inc()
		return nil, err
	}

	// Create snapshot directory
	snapDir := filepath.Join(s.root, "snapshots", key) // Use key instead of info.ID
	if err := os.MkdirAll(snapDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create snapshot dir: %w", err)
	}

	// If has parent, need to pull layer
	if parent != "" {
		layerDigest := info.Labels["containerd.io/snapshot/cri.layer-digest"]
		if layerDigest != "" {
			if err := s.prepareLayer(ctx, info, layerDigest, key); err != nil { // Pass key as id
				return nil, fmt.Errorf("failed to prepare layer: %w", err)
			}

			// Trigger prefetch if enabled
			if s.prefetch != nil {
				go s.prefetch.Prefetch(ctx, layerDigest, key) // Use key instead of info.ID
			}
		}
	}

	metrics.ActiveSnapshots.Inc()
	return s.mounts(info, key) // Pass key as id
}

func (s *Snapshotter) prepareLayer(ctx context.Context, info snapshots.Info, layerDigest string, id string) error {
	logrus.Infof("preparing layer %s for snapshot %s", layerDigest, id)

	// Pull layer content with concurrent downloading
	if err := s.puller.PullLayer(ctx, layerDigest, id); err != nil {
		return fmt.Errorf("failed to pull layer: %w", err)
	}

	return nil
}

func (s *Snapshotter) mounts(info snapshots.Info, id string) ([]mount.Mount, error) {
	snapDir := filepath.Join(s.root, "snapshots", id)
	
	if info.Kind == snapshots.KindActive {
		// Active snapshot: overlayfs with upper/work dirs
		lowerDir := filepath.Join(snapDir, "fs")
		upperDir := filepath.Join(snapDir, "upper")
		workDir := filepath.Join(snapDir, "work")
		
		for _, dir := range []string{lowerDir, upperDir, workDir} {
			if err := os.MkdirAll(dir, 0755); err != nil {
				return nil, err
			}
		}

		return []mount.Mount{
			{
				Type:   "overlay",
				Source: "overlay",
				Options: []string{
					fmt.Sprintf("lowerdir=%s", lowerDir),
					fmt.Sprintf("upperdir=%s", upperDir),
					fmt.Sprintf("workdir=%s", workDir),
				},
			},
		}, nil
	}

	// View/Committed: read-only bind mount
	return []mount.Mount{
		{
			Type:   "bind",
			Source: filepath.Join(snapDir, "fs"),
			Options: []string{"ro", "rbind"},
		},
	}, nil
}

func parseSize(s string) int64 {
	// Simplified size parsing
	if s == "" {
		return 10 * 1024 * 1024 * 1024 // 10GB default
	}
	// TODO: implement proper size parsing
	return 10 * 1024 * 1024 * 1024
}
